# _*_ coding: utf-8 _*_
# ---------------------------
# @创建作者：ZQQ
# @创建时间：2025/10/15
# @说明：MinIO代码和模型下载服务
# ---------------------------

import argparse
import os
import subprocess
import signal
import shutil
import threading
import time
import socket
import requests
from pathlib import Path

from flask import Flask, jsonify, request
from flask_cors import CORS
from utils.modelget import ModelDownloader
from utils.pyget import PyDownloader
from utils.compose_generator import DockerComposeGenerator
from utils.task_manager import TaskManager, TaskStatus, task_manager

app = Flask(__name__)
CORS(app)


def is_port_in_use(port):
    """
    检查端口是否已被占用
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(1)  # 设置超时时间
            result = s.connect_ex(('localhost', port))
            return result == 0
    except Exception as e:
        print(f"⚠️ 检查端口 {port} 时发生错误: {e}")
        return False


def send_deploy_info_to_get_server(deploy_info, notification_url=None):
    """
    向get_server.py发送部署信息
    """
    try:
        # 使用默认URL或传入的URL
        url = notification_url or "http://localhost:6010/received/"
        response = requests.post(url, json=deploy_info, timeout=5)
        print(f"📤 部署信息已发送到: {url} (状态码: {response.status_code})")
        return True
    except Exception as e:
        print(f"⚠️ 发送部署信息失败: {e}")
        return False


def async_deploy_task(task_id, code_url, model_url, model_name, port, service_folder, notification_url=None):
    """
    异步部署任务
    """
    task = task_manager.get_task(task_id)
    if not task:
        return
    
    try:
        # 步骤1: 下载代码文件
        task.update_status(TaskStatus.DOWNLOADING, "正在下载代码文件...", 10)
        code_downloader = PyDownloader(code_url, service_folder)
        code_path = code_downloader.download
        
        if not code_path:
            task.set_result(False, error="代码文件下载失败")
            return
        
        # 步骤2: 下载模型文件
        task.update_status(TaskStatus.DOWNLOADING, "正在下载模型文件...", 30)
        model_downloader = ModelDownloader(model_url, service_folder)
        model_path = model_downloader.download
        
        if not model_path:
            task.set_result(False, error="模型文件下载失败")
            return
        
        # 步骤3: 构建Docker镜像
        task.update_status(TaskStatus.BUILDING, "正在构建Docker镜像...", 50)
        success, message = build_and_start_service(service_folder, model_name, port)
        
        if not success:
            task.set_result(False, error=message)
            return
        
        # 步骤4: 启动服务
        task.update_status(TaskStatus.STARTING, "正在启动服务...", 80)
        time.sleep(5)  # 等待服务完全启动
        
        # 步骤5: 完成
        result = {
            'modelName': model_name,
            'serviceFolder': service_folder,
            'codePath': code_path,
            'modelPath': model_path,
            'port': port,
            'serviceUrl': f'http://localhost:{port}',
            'composeFile': f'docker-compose-{model_name}-{port}.yml'
        }
        task.set_result(True, result=result)
        
        # 步骤6: 发送部署信息到get_server.py
        deploy_info = {
            'taskId': task_id,
            'status': 'deploy_success',
            'modelName': model_name,
            'port': port,
            'serviceUrl': f'http://localhost:{port}',
            'serviceFolder': service_folder,
            'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
        }
        send_deploy_info_to_get_server(deploy_info, notification_url)
        
    except Exception as e:
        task.set_result(False, error=str(e))


def build_and_start_service(service_folder, model_name, port):
    """
    构建Docker镜像并启动服务
    """
    try:
        # 获取当前工作目录
        original_dir = Path.cwd()
        print(f"当前工作目录: {original_dir}")
        
        # 项目目录 - 使用新的文件夹结构
        project_dir = original_dir / service_folder
        
        if not project_dir.exists():
            return False, f"项目目录不存在: {project_dir}"
        
        print(f"📁 找到项目目录: {project_dir}")
        
        # 使用 compose_generator 生成带模型名称和端口号的 compose 文件
        print(f"🔧 生成 docker-compose 文件...")
        generator = DockerComposeGenerator()
        compose_path = generator.generate_for_project(
            folder_name=service_folder,
            model_name=model_name,
            port=port
        )
        
        print(f"📄 生成的 compose 文件: {compose_path}")
        
        # 停止可能存在的服务并清理容器
        print(f"🛑 停止可能存在的服务...")
        stop_cmd = f"docker-compose -f {compose_path} down"
        subprocess.run(stop_cmd, shell=True, capture_output=True)
        
        # 清理可能存在的同名容器（防止容器名称冲突）
        container_name = f"{model_name}-{port}-deploy"
        print(f"🧹 清理可能存在的容器: {container_name}")
        cleanup_cmd = f"docker rm -f {container_name} 2>/dev/null || true"
        subprocess.run(cleanup_cmd, shell=True, capture_output=True)
        
        # 使用docker-compose构建和启动服务
        print(f"🔨 开始使用docker-compose构建和启动服务...")
        
        # 设置端口环境变量
        env = os.environ.copy()
        env['APP_PORT'] = str(port)
        env['HOST_PORT_APP'] = str(port)
        env['MODEL_NAME'] = model_name
        
        # 构建和启动 - 使用 -f 参数指定文件路径
        compose_cmd = f"docker-compose -f {compose_path} up -d --build"
        result = subprocess.run(compose_cmd, shell=True, capture_output=True, text=True, env=env)
        
        if result.returncode != 0:
            return False, f"Docker Compose启动失败: {result.stderr}"
        
        print(f"✅ Docker Compose服务启动成功")
        print(f"🌐 服务地址: http://localhost:{port}")
        print(f"🤖 模型名称: {model_name}")
        
        return True, f"服务启动成功，访问地址: http://localhost:{port}"
        
    except Exception as e:
        return False, f"构建和启动服务失败: {str(e)}"


@app.route('/deploy/', methods=['POST'])
def deploy():
    """
    异步下载代码和模型文件，构建Docker镜像并启动服务
    请求参数:
    {
        "codeUrl": "代码文件的下载URL",
        "modelUrl": "模型文件的下载URL", 
        "modelName": "模型名称(必需)",
        "folderName": "存储文件夹名称(可选，默认为deploy)",
        "port": "服务端口(可选，默认为5000)",
        "notificationUrl": "部署通知URL(可选，默认为http://localhost:6010/received/)"
    }
    返回:
    {
        "success": true,
        "taskId": "任务ID",
        "message": "部署任务已启动",
        "statusUrl": "/status/{taskId}"
    }
    """
    try:
        data = request.get_json()
        
        # 获取参数
        code_url = data.get('codeUrl', '')
        model_url = data.get('modelUrl', '')
        model_name = data.get('modelName', '')
        folder_name = data.get('folderName', 'deploy')
        port = data.get('port', 5000)
        notification_url = data.get('notificationUrl', 'http://localhost:6010/received/')
        
        # 参数校验
        if not code_url or not model_url or not model_name:
            return jsonify({
                'success': False,
                'message': '缺少必要参数：codeUrl、modelUrl 和 modelName'
            }), 400
        
        # 检查端口是否已被占用
        if is_port_in_use(port):
            return jsonify({
                'success': False,
                'message': f'端口 {port} 已被占用，请选择其他端口或先停止占用该端口的服务'
            }), 400
        
        # 生成基于模型名称和端口号的文件夹名称
        service_folder = f"{model_name}-{port}"
        
        # 创建部署任务
        task_id = task_manager.create_task(model_name, port, service_folder)
        
        print(f"📥 开始异步部署任务...")
        print(f"🆔 任务ID: {task_id}")
        print(f"📁 服务文件夹: {service_folder}")
        print(f"🤖 模型名称: {model_name}")
        print(f"🔗 代码URL: {code_url}")
        print(f"🔗 模型URL: {model_url}")
        print(f"🌐 服务端口: {port}")
        print(f"✅ 端口 {port} 检查通过，未被占用")
        
        # 启动异步任务
        thread = threading.Thread(
            target=async_deploy_task,
            args=(task_id, code_url, model_url, model_name, port, service_folder, notification_url)
        )
        thread.daemon = True
        thread.start()
        
        return jsonify({
            'success': True,
            'taskId': task_id,
            'message': '部署任务已启动，请使用taskId查询状态',
            'statusUrl': f'/status/{task_id}',
            'data': {
                'modelName': model_name,
                'serviceFolder': service_folder,
                'port': port
            }
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'启动部署任务失败: {str(e)}'
        }), 500


@app.route('/status/<task_id>', methods=['GET'])
def get_task_status(task_id):
    """
    查询部署任务状态
    """
    try:
        task = task_manager.get_task(task_id)
        
        if not task:
            return jsonify({
                'success': False,
                'message': '任务不存在'
            }), 404
        
        response_data = task.to_dict()
        
        # 添加成功状态标识
        if task.status == TaskStatus.SUCCESS:
            response_data['success'] = True
        elif task.status == TaskStatus.FAILED:
            response_data['success'] = False
        else:
            response_data['success'] = None  # 进行中
        
        return jsonify({
            'success': True,
            'data': response_data
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'查询任务状态失败: {str(e)}'
        }), 500


@app.route('/tasks/', methods=['GET'])
def list_tasks():
    """
    列出所有部署任务
    """
    try:
        tasks = task_manager.list_tasks()
        
        return jsonify({
            'success': True,
            'data': {
                'tasks': tasks,
                'total': len(tasks)
            }
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'获取任务列表失败: {str(e)}'
        }), 500


@app.route('/stop/', methods=['POST'])
def stop():
    """
    停止Docker Compose服务
    请求参数:
    {
        "modelName": "模型名称(必需)",
        "port": "端口号(必需)"
    }
    """
    try:
        data = request.get_json()
        model_name = data.get('modelName', '')
        port = data.get('port', None)
        
        # 参数校验
        if not model_name or port is None:
            return jsonify({
                'success': False,
                'message': '缺少必要参数：modelName 和 port'
            }), 400
        
        # 生成服务文件夹名称
        service_folder = f"{model_name}-{port}"
        
        # 获取当前工作目录
        original_dir = Path.cwd()
        
        # 项目目录 - 使用新的文件夹结构
        project_dir = original_dir / service_folder
        
        if not project_dir.exists():
            return jsonify({
                'success': False,
                'message': f'项目目录不存在: {project_dir}'
            }), 400
        
        # 构建 compose 文件路径
        compose_filename = f"docker-compose-{model_name}-{port}.yml"
        compose_path = project_dir / compose_filename
        
        if not compose_path.exists():
            return jsonify({
                'success': False,
                'message': f'Docker Compose文件不存在: {compose_path}'
            }), 400
        
        print(f"🛑 停止Docker Compose服务: {compose_path}")
        print(f"📁 服务文件夹: {service_folder}")
        print(f"🤖 模型名称: {model_name}")
        print(f"🌐 端口号: {port}")
        
        # 使用docker-compose停止服务 - 使用 -f 参数
        stop_cmd = f"docker-compose -f {compose_path} down"
        result = subprocess.run(stop_cmd, shell=True, capture_output=True, text=True)
        
        if result.returncode != 0:
            return jsonify({
                'success': False,
                'message': f'停止服务失败: {result.stderr}'
            }), 400
        
        # 停止服务成功
        print(f"✅ Docker Compose服务已停止")
        stop_message = f"Docker Compose服务已停止 (模型: {model_name}, 端口: {port})"
        
        return jsonify({
            'success': True,
            'message': stop_message
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'停止服务失败: {str(e)}'
        }), 500


@app.route('/delete/', methods=['POST'])
def delete():
    """
    删除Docker容器、镜像并清理项目文件夹
    请求参数:
    {
        "modelName": "模型名称(必需)",
        "port": "端口号(必需)"
    }
    """
    try:
        data = request.get_json()
        model_name = data.get('modelName', '')
        port = data.get('port', None)
        
        # 参数校验
        if not model_name or port is None:
            return jsonify({
                'success': False,
                'message': '缺少必要参数：modelName 和 port'
            }), 400
        
        # 生成服务文件夹名称
        service_folder = f"{model_name}-{port}"
        
        # 获取当前工作目录
        original_dir = Path.cwd()
        
        # 项目目录 - 使用新的文件夹结构
        project_dir = original_dir / service_folder
        
        if not project_dir.exists():
            return jsonify({
                'success': False,
                'message': f'项目目录不存在: {project_dir}'
            }), 400
        
        # 构建 compose 文件路径
        compose_filename = f"docker-compose-{model_name}-{port}.yml"
        compose_path = project_dir / compose_filename
        
        if not compose_path.exists():
            return jsonify({
                'success': False,
                'message': f'Docker Compose文件不存在: {compose_path}'
            }), 400
        
        print(f"🗑️ 开始删除服务: {compose_path}")
        print(f"📁 服务文件夹: {service_folder}")
        print(f"🤖 模型名称: {model_name}")
        print(f"🌐 端口号: {port}")
        
        # 步骤1: 停止并删除Docker容器
        print(f"🛑 停止并删除Docker容器...")
        stop_cmd = f"docker-compose -f {compose_path} down"
        result = subprocess.run(stop_cmd, shell=True, capture_output=True, text=True)
        
        if result.returncode != 0:
            print(f"⚠️ 停止容器失败: {result.stderr}")
        
        # 步骤2: 删除Docker镜像
        print(f"🗑️ 删除Docker镜像...")
        image_name = f"{model_name}-{port}"
        delete_image_cmd = f"docker rmi -f {image_name} 2>/dev/null || true"
        subprocess.run(delete_image_cmd, shell=True, capture_output=True, text=True)
        
        # 步骤3: 清理项目文件夹（包括下载的代码和模型文件）
        print(f"🧹 清理项目文件夹: {project_dir}")
        try:
            shutil.rmtree(project_dir)
            print(f"✅ 项目文件夹已清理: {project_dir}")
            cleanup_message = f"服务已完全删除 (模型: {model_name}, 端口: {port})"
        except Exception as cleanup_error:
            print(f"⚠️ 清理项目文件夹失败: {cleanup_error}")
            cleanup_message = f"容器和镜像已删除，但清理项目文件夹失败: {cleanup_error} (模型: {model_name}, 端口: {port})"
        
        return jsonify({
            'success': True,
            'message': cleanup_message
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'删除服务失败: {str(e)}'
        }), 500


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='MinIO下载服务')
    parser.add_argument('--port', type=int, default=8050, help='服务端口')
    args = parser.parse_args()
    
    app.run(host='0.0.0.0', port=args.port, debug=False)